package ws.very.util.dingtalk.encryptor

import java.io.ByteArrayOutputStream
import java.nio.charset.Charset
import java.security.MessageDigest
import java.util.Arrays
import javax.crypto.Cipher
import javax.crypto.spec.{IvParameterSpec, SecretKeySpec}

import org.apache.commons.codec.binary.Base64
import ws.very.util.dingtalk.{PKCS7Padding, Utils}

import scala.util.{Failure, Success, Try}

/**
  * 构造函数
  *
  * @param token          钉钉开放平台上，开发者设置的token
  * @param encodingAesKey 钉钉开放台上，开发者设置的EncodingAESKey
  * @param corpId         ISV进行配置的时候应该传对应套件的SUITE_KEY，普通企业是Corpid
  */
class DingTalkEncryptor(token: String, encodingAesKey: String, corpId: String) {

  import DingTalkEncryptException._

  private val CHARSET: Charset = Charset.forName("utf-8")
  private val base64: Base64 = new Base64
  /** ask getPaddingBytes key固定长度 **/
  private val AES_ENCODE_KEY_LENGTH: Int = 43
  /** 加密随机字符串字节长度 **/
  private val RANDOM_LENGTH: Int = 16

  private val aesKey = Base64.decodeBase64(encodingAesKey + "=")


  assert(encodingAesKey.length == AES_ENCODE_KEY_LENGTH, getMsg(AES_KEY_ILLEGAL))

  /**
    * 将和钉钉开放平台同步的消息体加密,返回加密Map
    *
    * @param plaintext 传递的消息体明文
    * @param timeStamp 时间戳
    * @param nonce     随机字符串
    * @return
    */
  def getEncryptedMap(plaintext: String, timeStamp: String, nonce: String) = {
    // 加密
    encrypt(Utils.getRandomStr(RANDOM_LENGTH), plaintext).flatMap { encryptString =>
      getSignature(token, timeStamp, nonce, encryptString).map(_ -> encryptString)
    } map { case (signature, encryptString) =>
      Map(
        "msg_signature" -> signature,
        "encrypt" -> encryptString,
        "timeStamp" -> timeStamp,
        "nonce" -> nonce
      )
    }
  }

  /**
    * 密文解密
    *
    * @param msgSignature 签名串
    * @param timeStamp    时间戳
    * @param nonce        随机串
    * @param encryptMsg   密文
    * @return 解密后的原文
    */
  def getDecryptMsg(msgSignature: String, timeStamp: String, nonce: String, encryptMsg: String) = {
    //校验签名
    getSignature(token, timeStamp, nonce, encryptMsg).flatMap { signature =>
      if (signature != msgSignature) {
        Failure(DingTalkEncryptException(COMPUTE_SIGNATURE_ERROR))
      } else {
        // 解密
        decrypt(encryptMsg)
      }
    }
  }

  /*
   * 对明文加密.
   * @param text 需要加密的明文
   * @return 加密后base64编码的字符串
   */
  private def encrypt(random: String, plaintext: String) = {
    val randomBytes: Array[Byte] = random.getBytes(CHARSET)
    val plainTextBytes: Array[Byte] = plaintext.getBytes(CHARSET)
    val lengthByte: Array[Byte] = Utils.int2Bytes(plainTextBytes.length)
    val corpidBytes: Array[Byte] = corpId.getBytes(CHARSET)
    val byteStream: ByteArrayOutputStream = new ByteArrayOutputStream
    byteStream.write(randomBytes)
    byteStream.write(lengthByte)
    byteStream.write(plainTextBytes)
    byteStream.write(corpidBytes)
    val padBytes: Array[Byte] = PKCS7Padding.getPaddingBytes(byteStream.size)
    byteStream.write(padBytes)
    val unencrypted: Array[Byte] = byteStream.toByteArray
    byteStream.close()
    Try {
      val cipher: Cipher = Cipher.getInstance("AES/CBC/NoPadding")
      val keySpec: SecretKeySpec = new SecretKeySpec(aesKey, "AES")
      val iv: IvParameterSpec = new IvParameterSpec(aesKey, 0, 16)
      cipher.init(Cipher.ENCRYPT_MODE, keySpec, iv)
      val encrypted: Array[Byte] = cipher.doFinal(unencrypted)
      base64.encodeToString(encrypted)
    } match {
      case Failure(_) => Failure(DingTalkEncryptException(DingTalkEncryptException.COMPUTE_ENCRYPT_TEXT_ERROR))
      case a => a
    }
  }


  /*
   * 对密文进行解密.
   * @param text 需要解密的密文
   * @return 解密得到的明文
   */
  private def decrypt(text: String) = {

    val originalArr = Try {
      // 设置解密模式为AES的CBC模式
      val cipher: Cipher = Cipher.getInstance("AES/CBC/NoPadding")
      val keySpec: SecretKeySpec = new SecretKeySpec(aesKey, "AES")
      val iv: IvParameterSpec = new IvParameterSpec(Arrays.copyOfRange(aesKey, 0, 16))
      cipher.init(Cipher.DECRYPT_MODE, keySpec, iv)
      // 使用BASE64对密文进行解码
      val encrypted: Array[Byte] = Base64.decodeBase64(text)
      // 解密
      cipher.doFinal(encrypted)
    } match {
      case e: Failure[_] =>
        Failure(DingTalkEncryptException(DingTalkEncryptException.COMPUTE_DECRYPT_TEXT_ERROR))
      case a => a
    }
    originalArr flatMap { originalArr =>
      val bytes: Array[Byte] = PKCS7Padding.removePaddingBytes(originalArr)
      val networkOrder: Array[Byte] = Arrays.copyOfRange(bytes, 16, 20)
      val plainTextLegth: Int = Utils.bytes2int(networkOrder)
      val plainText = new String(Arrays.copyOfRange(bytes, 20, 20 + plainTextLegth), CHARSET)
      val fromCorpid = new String(Arrays.copyOfRange(bytes, 20 + plainTextLegth, bytes.length), CHARSET)
      if (fromCorpid != corpId) {
        println(s"==>${fromCorpid}",corpId)
        Failure(DingTalkEncryptException(COMPUTE_DECRYPT_TEXT_CORPID_ERROR))
      } else {
        Success(plainText)
      }
    }
  }

  /**
    * 数字签名
    *
    * @param token     isv token
    * @param timestamp 时间戳
    * @param nonce     随机串
    * @param encrypt   加密文本
    * @return
    * DingTalkEncryptException
    */
  def getSignature(token: String, timestamp: String, nonce: String, encrypt: String) = {
    val str = Array[String](token, timestamp, nonce, encrypt).sorted.mkString("")
    Try {
      val md: MessageDigest = MessageDigest.getInstance("SHA-1")
      md.update(str.getBytes)
      val digest: Array[Byte] = md.digest
      val hexstr: StringBuffer = new StringBuffer
      var shaHex: String = ""
      var i: Int = 0
      while (i < digest.length) {
        shaHex = Integer.toHexString(digest(i) & 0xFF)
        if (shaHex.length < 2) hexstr.append(0)
        hexstr.append(shaHex)
        i += 1
      }
      hexstr.toString
    } match {
      case Failure(_) => Failure(DingTalkEncryptException(COMPUTE_SIGNATURE_ERROR))
      case a => a
    }
  }

}
